// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <melon/executors/cpu_thread_pool_executor.h>
#include <pollux/common/base/semaphore.h>

using namespace kumo::pollux;
std::atomic<int32_t> numReleased = 0;

struct Callable {
  Callable(Semaphore& sem) : sem(sem) {}
  Semaphore& sem;

  void operator()() {
    ++numReleased;
    sem.release();
  }
};

TEST(SemaphoreTest, threads) {
  // Makes a pool of producer threads that release a semaphore and a set of
  // consumer threads that acquire the same semaphore. Once a consumer sees that
  // the expected number of acquires have been done, it releases the semaphore
  // enough times to unblock the other consumers.
  constexpr int32_t kNumProducers = 20;
  constexpr int32_t kNumConsumers = 20;
  constexpr int32_t kNumOps = 10000;
  auto executor = std::make_unique<melon::CPUThreadPoolExecutor>(kNumProducers);
  Semaphore sem(0);
  std::atomic<int32_t> numDone = 0;
  std::vector<std::thread> consumers;
  consumers.reserve(kNumConsumers);
  for (auto i = 0; i < kNumConsumers; ++i) {
    consumers.emplace_back(std::thread([&]() {
      for (;;) {
        sem.acquire();
        int32_t done = ++numDone;
        if (numDone == kNumOps) {
          // All producers are finished, continue the other consumers.
          for (auto i = 0; i < kNumConsumers - 1; ++i) {
            ++numReleased;
            sem.release();
          }
        }
        if (done >= kNumOps) {
          return;
        }
      }
    }));
  }
  std::vector<Callable> ops;
  ops.reserve(kNumOps);
  for (auto i = 0; i < kNumOps; ++i) {
    ops.emplace_back(sem);
    executor->add(ops.back());
  }
  for (auto& thread : consumers) {
    thread.join();
  }
  EXPECT_EQ(kNumOps + kNumConsumers - 1, numDone);
  EXPECT_EQ(numDone, numReleased);
}
